"""PlotGrid, PlotStack, NavalaPlot, DailyAxis.

_* methods and variables are called only among the classes defined within this
module.
"""
from PyQt4.QtCore import *
import numpy as np
import datetime
import pyqtgraph as gr
import math
import logging

from misc import misc
from misc import timedata
from block import block

################################################################################
## PlotGrid

class PlotGrid(gr.GraphicsLayout):
    """Collection of PlotStack objects sharing same time limits."""
    
    # TODO: maybe will need to flag need to re-setup and override paint()
    # to automatically call Setup()

    # Emits this signal whenever the grid wants to inform that there is new
    # information about a new mouse position
    PostCrosshairTuples = pyqtSignal(list)
    
    # Quote period. Initially set to -1
    period = -1
    
    @property
    def stacks(self):
        """Getter for the stacks=(list of PlotStack) property."""
        return self._stacks

    def __init__(self, *args, **kwargs):
        gr.GraphicsLayout.__init__(self, *args, **kwargs)
        self.layout.setSpacing(2) # both horizontal and vertical, between plots
        self.layout.setContentsMargins(2, 2, 2, 2)
        self.__numItems = -1
        self.__nr = -1
        self.__nc = -1
        self.__crosshairPlot = None

        # PlotStack list
        self._stacks = []
        # TimeData list
        self._dataDict = []
        # Set to True whenever new PlotStack is added
        self._flagNeedsSetup = False
        # Set to True whenever new TimeData is registered. Currently not
        # checked upon, just being set.
        self._flagNeedsUpdate = False

    def SetNumStacks(self, n):
        """This method must be called before items start to be added.
       
        This is because the number of rows (and more importantly, columns)
        needs to be calculated first."""
        self.__numItems = n
        (self.__nr, self.__nc) = self.__CalculateDims(n)
        # unnecessary self.__row = 0
        self.__col = 0
        self._stacks = [] # Separate because the original items is a dictionary, and I need some ORDER!
###        self._flagNeedsSetup = True

    def SetTimeData(self, dd):
        """Each element in dd is passed to one item in self.items"""
        for (stack, d) in zip(self._stacks, dd):
            stack._SetTimeData(d)
        self._flagNeedsUpdate = True
        
    def Update(self):
        """Must be called after data assignment process.
        
        Tasks include:
            - timedata.TrimTimeGaps()
            - timedata.IntersectTimestamps()
            - calling _Update() for all the items in _stacks."""
        
        dd = self._dataDict
        if len(dd) > 0:
            timedata.TrimTimeGaps(dd, self.period) # This will set timestamp2
            # timestamps original, timesTamps warped
            (tto, ttw) = timedata.IntersectTimestamps(dd)
            if len(tto) > 0:
                tMin = min(ttw)
                tMax = max(ttw)

            for (i, stack) in enumerate(self._stacks):
                if len(tto) > 0:
                    stack._SetXRange([tMin, tMax])
                    stack._SetTimestamps(tto, ttw)
                else:
                    stack._SetXRange([0, 1])
        for obj in self._stacks:
            obj._Update()
        self._flagNeedsUpdate = False
#        logging.debug("=========>>> PlotGrid().Update() has been called")


    ### --- INTERNAL --- #######################################################

    def _Setup(self):
        """GUI-related setup. Updates state after _stacks is changed.
        
        - Automatically handled
        - Calls _Setup for * elements in _stacks."""

        for (i, stack) in enumerate(self._stacks):
            if self._GetHasBottomAxis(i):
                pass
#                stack.SetBottomTicks(tt)
            else:
                stack._HideBottomAxis()
        
        for obj in self._stacks:
            obj._Setup()
            
        self._flagNeedsSetup = False
            
    def _GetHasBottomAxis(self, idx):
        """Gets whether plot as position idx has bottom axis.
        
        Returns True for the items that that are the last of their column."""
        c = idx % self.__nc
        r = idx/self.__nc
        flag = (r == self.__nr-1) or r == self.__nr-2 and c > (self.__numItems-1)%self.__nc
        return flag

    def __CalculateDims(self, n):
        """Calculates number of rows and columns based on len(grid).

        Number of columns tends to be higher.
        
        Returns: (number of rows, number of columns)
        """
        if n == 0:
            return (0, 0)
        nc = int(math.floor(math.sqrt(n)))
        nr = int(math.ceil(float(n)/nc))
        return (nr, nc)

    def _AddStack(self, obj):
        """This is called from PlotStack.__init__().
        
        Adds PlotStack object and increments current cell (row+col) pos."""
        if self.__numItems == -1:
            raise RuntimeError("Please call SetNumStacks() before adding items!")
        if self.__col >= self.__nc:
            self.nextRow()
            self.__col = 0
        self.addItem(obj)
        self._stacks.append(obj)
        self.__col += 1
        self._flagNeedsSetup = True
        
    def _RegisterTimeData(self, td):
        if not td in self._dataDict:
            self._dataDict.append(td)
            self._flagNeedsUpdate = True

    def _UnregisterTimeData(self, td):
        """Opposite of _RegisterTimeData()"""
        if td in self._dataDict:
            self._dataDict.remove(td)
            self._flagNeedsUpdate = True
        
    def _PostCrosshairPlot(self, plot):
        """Informs the grid which plot must have the crosshair."""
        if plot != self.__crosshairPlot:
            if self.__crosshairPlot is not None:
                self.__crosshairPlot._HideCrosshair()
            self.__crosshairPlot = plot
#            self.crosshairPlot._ShowCrosshair()

    def HideCrosshair(self):
        """TODO: why this does not start with an underscore?"""
        if self.__crosshairPlot is not None:
            self.__crosshairPlot._HideCrosshair()

    def paint(self, *args, **kwargs):
        """paint() method overriden to check for _flagNeedsSetup."""
        
# (3/June/2013) This haas been apparently sorted, let's leave it for a bit more    
#########        if self._flagNeedsUpdate:
#########            # Protection against not calling Update()
#########            logging.critical("W W W W W W W W W T T T T T T T T F F F F F F F F !!! PlotGrid.paint() called before Update()")
#########        else:
#########            logging.debug("Cooooooooooool, Update() has been called before paint()")
        if self._flagNeedsSetup:
            self._Setup()
        gr.GraphicsLayout.paint(self, *args, **kwargs)

################################################################################
## PlotStack

class PlotStack(gr.GraphicsLayout):
    """Non-visual representation of a collection of pyqtgraph.PlotItem objects"""
    
    @property
    def plots(self):
        """Getter for the plots property (read-only)."""
        return self._plots
    
    def __init__(self, plotGrid, sTitle="Untitled", *args, **kwargs):
        gr.GraphicsLayout.__init__(self, *args, **kwargs)
        # List of NavalaPlot objects
        self._plots = []
        # title label
        self._titleLabel = None
        self._grid = plotGrid
        plotGrid._AddStack(self)
        self.layout.setContentsMargins(1, 1, 1, 1)
        l = gr.LabelItem(sTitle)
        self._titleLabel = l
        self.addItem(l)
        self.layout.setSpacing(2)
        self.layout.setRowStretchFactor(0, 0) # Label will takes minimum height
        self.nextRow()


    def _SetTimeData(self, td):
        """Called from PlotGrid.SetTimeData().
        
        This routine mustn't be called directly."""
#        self.timeData = timeData
        for obj in self._plots:
            obj._SetTimeData(td)

    def _AddPlot(self, plot):
        """This is called from NavalaPlot.__init__()."""
        self.addItem(plot)
        self._plots.append(plot)

    def _RegisterTimeData(self, td):
        """This is called by a Plot whenever it gets a new TD from a Drawer."""
        self._grid._RegisterTimeData(td)

    def _UnregisterTimeData(self, td):
        """Opposite of _RegisterTimeData()"""
        self._grid._UnregisterTimeData(td)

    def _HideBottomAxis(self):
        for item in self.items.iterkeys():
            if hasattr(item, "hideAxis"):
                item.hideAxis("bottom")

    def _SetTimestamps(self, tto, ttw):
        a = self._plots[-1].getAxis("bottom")
        a.SetTimestamps(tto, ttw)
        #self.prepareGeometryChange()

    def _SetXRange(self, range):
        for item in self.items.iterkeys():
            if hasattr(item, "setRange"):
                pass
                #TODO this is left without a choice for now
                item.setRange(xRange=range)
                
    def _Setup(self):
        """Called after all plots are added. Propagates to all plots."""
        for obj in self._plots:
            obj._Setup()

    def _Update(self):
        """Called after all data assignment. Propagates to all plots."""
        for obj in self._plots:
            obj._Update()

    def _PostCrosshairPlot(self, plot):
        """Passes on to the PlotGrid."""
        self._grid._PostCrosshairPlot(plot)

    def _PostCrosshairX(self, x, t):
        """Talks to all Drawers in stack to produce their coordinates.
        
        Arguments:
            x -- x-mouse coordinate translated into the scale of timestamp2
                 of the reference TimeData of one of the plots in the stack.
            t -- a timestamp value (not timestamp2) matching to x"""
        a = [("time", misc.ts2dt(t).strftime("%d/%b/%Y %H:%M:%S"))]
        for plot in self._plots:
            a += plot._CollectCrosshairTuples(x)
        self._grid.PostCrosshairTuples.emit(a)



################################################################################
## Plots

class NavalaPlot(gr.PlotItem):
    @property
    def drawers(self):
        """Getter for the drawers property (read only)."""
        return self._drawers
    
    # (Read-only)
    @property
    def dataBlocks(self):
        return self._dataBlocks
    

    def AddDataBlock(self, name):
        """Adds new PlotDataBlock and returns it."""
        temp = PlotDataBlock()
        self._dataBlocks[name] = temp
        return temp

    def AddCurve(self, curve):
        self._curves.append(curve)
        self.addItem(curve)

    def __init__(self, plotStack, *args, **kwargs):
        ax = DailyAxis("bottom")
        gr.PlotItem.__init__(self, axisItems={"bottom": ax}, border=(192, 192, 192), *args, **kwargs)
        self._stack = plotStack
        plotStack._AddPlot(self)
        
        #crosshair
        self.vLine = gr.InfiniteLine(angle=90, movable=False)
        self.hLine = gr.InfiniteLine(angle=0, movable=False)
        self.addItem(self.hLine, ignoreBounds=True)
        self.addItem(self.vLine, ignoreBounds=True)
        self.hLine.setVisible(False)
        self.vLine.setVisible(False)        

        # List of curves
        self._curves = []
        # List of Drawer blocks
        self._drawers = []
        # List of Data blocks
        self._dataBlocks = {}
        # Reference timedata object for the Crosshair search
        self._refTd = None
        # Not Empty Data flag
        self._flagHasData = False

    def _AddDrawer(self, obj):
        """This mustn't be called directly, auto from Drawer.__init__()."""
        self._drawers.append(obj)
        
    def _RegisterTimeData(self, td):
        """This is called by a Drawer whenever a new TimeData is set."""
        if self._refTd is None:
            self._refTd = td
            self._flagHasData = len(td) > 0
        self._stack._RegisterTimeData(td)
        
    def _UnregisterTimeData(self, td):
        """Opposite of _RegisterTimeData()"""
        if td == self._refTd:
            self._refTd = None
            # Does not check if there could be another td for reference,
            # but doesn't matter really
            self._flagHasData = False
        self._stack._UnregisterTimeData(td)

    def _SetTimeData(self, td):
        """Calls SetTimeData() for all dataBlocks.

        Caution: use either _SetTimeData() or _RegisterTimeData(). The former
        when all drawers in the plot use same TimeData; the later when there
        are multiple TimeData sources in the plot. Using those both methods
        will cause the TimeData list in the PlotGrid to become messed up."""
#######################        if self._refTd is not None:
#######################            self._UnregisterTimeData(self._refTd)
#######################        self._RegisterTimeData(td)
        if len(self._dataBlocks) == 0:
            raise RuntimeError("%s object has no data blocks!" % (self.__class__.__name__))
        for obj in self._dataBlocks.itervalues():
            obj.SetTimeData(td)
                
# I think the bounding rect of the plot won't change        self.prepareGeometryChange()
#        self.update()

    def _CollectCrosshairTuples(self, x):
        """Collects coordinate information from all drawers into a list."""
        a = []
        for drawer in self._drawers:
            a += drawer.GetCrosshairTuples(x)
        return a

    def mouseMoved(self, evt):
        if not self._flagHasData:
            return
        pos = evt[0]  ## using signal proxy turns original arguments into a tuple
        if self.sceneBoundingRect().contains(pos):
            mousePoint = self.vb.mapSceneToView(pos)
            index = misc.BSearchRound(self._refTd.timestamp2, mousePoint.x())
            self._crosshairX = self._refTd.timestamp2[index]
            self._crosshairY = mousePoint.y() 
            self._ShowCrosshair()
            self._stack._PostCrosshairPlot(self)
            # Sends mouse-x converted to timestamp2 scale to be searched by
            # drawers inside their respective timestamp2's; and also
            # the timestamp (correct time) value because the PlotStack, does not
            # have/(need to have) a reference TimeData (and also reduces the need
            # for 1 binary search)
            t = self._refTd.timestamp[index]
            self._stack._PostCrosshairX(mousePoint.x(), t)

    def _ShowCrosshair(self):
        self.vLine.setPos(self._crosshairX)
        self.hLine.setPos(self._crosshairY)
        self.hLine.setVisible(True)
        self.vLine.setVisible(True)
        
    def _HideCrosshair(self):
        self.hLine.setVisible(False)
        self.vLine.setVisible(False)        

    def _Setup(self):
        """This is called after the plot has a scene."""
        self.proxy = gr.SignalProxy(self.scene().sigMouseMoved, 
                                    rateLimit=5, # signals per second
                                    slot=self.mouseMoved)

    def _Update(self):
        """This is called after a data assignment process.
        
        Updates plot y-range.
        """
        for drawer in self._drawers:
            drawer._Update()
        self.getViewBox().updateAutoRange()
        for curve in self._curves:
            curve.Update()
            
            


################################################################################

class DailyAxis(gr.AxisItem):
    timestamp = []
    timestamp2 = []
    def tickStrings(self, values2, scale, spacing):
        """Context formating of date information.
        For example, if only the day changes, there is no need to include the
        year and month in the tick string.
        """
        #print "ASKING FOR TICKSTRINGS..."
        if len(self.timestamp) == 0:
            return gr.AxisItem.tickStrings(self, values2, scale, spacing)
            
        strs = []
        values = [self.timestamp[misc.BSearchRound(self.timestamp2, value2)]
                   for value2 in values2]
        #if self.gapRemaker:
        #    values = self.gapRemaker.Convert(values)
        dts = [misc.ts2dt(x) for x in values]
        bFirst = True
        for dt in dts:
            if bFirst:
                fmt = "%d/%b/%Y" # First tick needs to inform the complete time location
                bFirst = False
            else:
                if dt.year != y: # Year changed
                    fmt = "%d/%b/%Y"
                elif dt.month != m: # ...
                    fmt = "%d/%b/%y"
                elif dt.day != d:
                    fmt = "%d"
                else:
                    fmt ="%H:%M"
                
            strs.append(dt.strftime(fmt))
            y, m, d = (dt.year, dt.month, dt.day)        

        return strs
        
    #def SetGapRemaker(self, obj):
    #    self.gapRemaker = obj
     
    def SetTimestamps(self, tss, tss2):
        self.timestamp = tss
        self.timestamp2 = tss2
        
    def tickSpacing(self, minVal, maxVal, size):
        """TODO: I should be able to calculate the precise tick locations instead of using this.
        Actually I think that I need a lower-level graphics API."""
        return [gr.AxisItem.tickSpacing(self, minVal, maxVal, size)[0]]


################################################################################

class PlotDataBlock(block.OBlock):
    """Associates with a single TimeData registered within a plot.

    TODO: move this to a different file so that all block classes stay within the "block" directory."""
    
    def SetTimeData(self, td):
        """Directs data to output."""
        self["output"][0] = td
